import concurrent.futures
import threading

from utils.logger import log


class TaskState:
    RUNNING = "运行中"
    DONE = "完成"
    CANCEL = "取消"
    ERROR = "错误"
    WAITING = "等待"
    UNKNOWN = "未知"

class TaskManager:
    def __init__(self, max_workers=5):
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
        self.tasks = {}  # task_id -> {'future': ..., 'stop_event': ...}

    def start_task(self, task_id, task_fn, *args, **kwargs):
        if task_id in self.tasks:
            msg = f"任务 {task_id} 已存在"
            log.error(msg)
            return False, msg

        stop_event = threading.Event()
        future = self.executor.submit(task_fn, stop_event, *args, **kwargs)
        log.info("已经启动", extra={"pk": task_id})
        self.tasks[task_id] = {'future': future, 'stop_event': stop_event}

        return True, None

    def stop_task(self, task_id):
        task = self.tasks.get(task_id)
        if task:
            task['stop_event'].set()
            log.info("发出停止信号，稍等...", extra={"pk": task_id})

        return self.get_status(task_id)


    def get_status(self, task_id):
        task = self.tasks.get(task_id)
        if not task:
            return TaskState.UNKNOWN
        future = task['future']
        if future.cancelled():
            return TaskState.CANCEL
        elif future.running():
            return TaskState.RUNNING
        elif future.done():
            error = future.exception()
            if error:
                self.tasks[task_id]['error'] = error
                return TaskState.ERROR
            else:
                return TaskState.DONE
        else:
            return TaskState.WAITING

    def get_all_statuses(self):
        return {tid: self.get_status(tid) for tid in self.tasks}

    def stop_all_tasks(self):
        for task_id, info in self.tasks.items():
            info['stop_event'].set()

    def shutdown(self):
        self.stop_all_tasks()
        # self.executor.shutdown(wait=True)


# 模块级单例
task_manager = TaskManager()
